{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Notebook to demonstrate Auto-Labeling workflow\n",
    "\n",
    "Transfer learning is the process of transferring learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task. Train Adapt Optimize (TAO) Toolkit  is a simple and easy-to-use Python based AI toolkit for taking purpose-built AI models and customizing them with users' own data.\n",
    "\n",
    "![image](https://d29g4g2dyqv443.cloudfront.net/sites/default/files/akamai/TAO/tlt-tao-toolkit-bring-your-own-model-diagram.png)\n",
    "\n",
    "\n",
    "### The workflow in a nutshell\n",
    "\n",
    "- Creating a dataset\n",
    "- Upload dataset to the service\n",
    "- Getting a PTM from NGC\n",
    "- Model Actions\n",
    "    - Train (Normal/AutoML)\n",
    "    - Evaluate\n",
    "    - Inference on TAO\n",
    "\n",
    "### Table of contents\n",
    "\n",
    "1. [Create datasets ](#head-1)\n",
    "1. [List the created datasets](#head-2)\n",
    "1. [Create an experiment](#head-4)\n",
    "1. [List experiments](#head-5)\n",
    "1. [Assign train, eval datasets](#head-6)\n",
    "1. [Assign PTM](#head-7)\n",
    "1. [View hyperparameters that are enabled by default](#head-8)\n",
    "1. [Set AutoML related configurations](#head-9)\n",
    "1. [Actions](#head-10)\n",
    "1. [Train](#head-11)\n",
    "1. [Evaluate](#head-12)\n",
    "1. [TAO inference](#head-20)\n",
    "\n",
    "### Requirements\n",
    "Please find the server requirements [here](https://docs.nvidia.com/tao/tao-toolkit/text/tao_toolkit_api/api_setup.html#)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "import requests\n",
    "import glob\n",
    "import time\n",
    "from IPython.display import clear_output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### FIXME\n",
    "\n",
    "1. Assign a workdir in FIXME 1\n",
    "1. Assign the ip_address and port_number in FIXME 2 ([info](https://docs.nvidia.com/tao/tao-toolkit/text/tao_toolkit_api/api_rest_api.html))\n",
    "1. Assign the ngc_api_key variable in FIXME 3\n",
    "1. (Optional) Enable AutoML if needed in FIXME 4\n",
    "1. (Optional) Choose between bayesian and hyperband automl_algorithm in FIXME 5 (If automl was enabled in FIXME4)\n",
    "1. Choose to download jobs or not in FIXME 6\n",
    "1. Choose between default and custom dataset in FIXME 7\n",
    "1. Assign path of DATA_DIR in FIXME 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"mal\"\n",
    "workdir = \"workdir_auto_labeling\" # FIXME1\n",
    "host_url = \"http://<ip_address>:<port_number>\" # FIXME2 example: https://10.137.149.22:32334\n",
    "# In host machine, node ip_address and port number can be obtained as follows,\n",
    "# ip_address: hostname -i\n",
    "# port_number: kubectl get service ingress-nginx-controller -o jsonpath='{.spec.ports[0].nodePort}'\n",
    "ngc_api_key = \"<ngc_api_key>\" # FIXME3 example: (Add NGC API key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "automl_enabled = False # FIXME4 set to TRUE if you want to run automl for the model chosen in the previous cell\n",
    "automl_algorithm = \"bayesian\" # FIXME5 example: bayesian/hyperband\n",
    "# FIXME6 Defaulted to False as downloading jobs from service to your machine takes time\n",
    "# Set to True if you want to download jobs where examples have been provided like for train, export, inference.\n",
    "download_jobs = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Exchange NGC_API_KEY for JWT\n",
    "data = json.dumps({\"ngc_api_key\": ngc_api_key})\n",
    "response = requests.post(f\"{host_url}/api/v1/login\", data=data)\n",
    "user_id = response.json()[\"user_id\"]\n",
    "print(\"User ID\",user_id)\n",
    "token = response.json()[\"token\"]\n",
    "print(\"JWT\",token)\n",
    "\n",
    "# Set base URL\n",
    "base_url = f\"{host_url}/api/v1/users/{user_id}\"\n",
    "print(\"API Calls will be forwarded to\",base_url)\n",
    "\n",
    "headers = {\"Authorization\": f\"Bearer {token}\"}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creating workdir\n",
    "if not os.path.isdir(workdir):\n",
    "    os.makedirs(workdir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function to split tar files <a class=\"anchor\" id=\"head-1.1\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import tarfile\n",
    "\n",
    "def split_tar_file(input_tar_path, output_dir, max_split_size=0.2*1024*1024*1024):\n",
    "\tos.makedirs(output_dir, exist_ok=True)\n",
    "\t\n",
    "\twith tarfile.open(input_tar_path, 'r') as original_tar:\n",
    "\t\tmembers = original_tar.getmembers()\n",
    "\t\tcurrent_split_size = 0\n",
    "\t\tcurrent_split_number = 0\n",
    "\t\tcurrent_split_name = os.path.join(output_dir, f'smaller_file_{current_split_number}.tar')\n",
    "\t\t\n",
    "\t\twith tarfile.open(current_split_name, 'w') as split_tar:\n",
    "\t\t\tfor member in members:\n",
    "\t\t\t\tif current_split_size + member.size <= max_split_size:\n",
    "\t\t\t\t\tsplit_tar.addfile(member, original_tar.extractfile(member))\n",
    "\t\t\t\t\tcurrent_split_size += member.size\n",
    "\t\t\t\telse:\n",
    "\t\t\t\t\tsplit_tar.close()\n",
    "\t\t\t\t\tcurrent_split_number += 1\n",
    "\t\t\t\t\tcurrent_split_name = os.path.join(output_dir, f'smaller_file_{current_split_number}.tar')\n",
    "\t\t\t\t\tcurrent_split_size = 0\n",
    "\t\t\t\t\tsplit_tar = tarfile.open(current_split_name, 'w')  # Open a new split tar archive\n",
    "\t\t\t\t\tsplit_tar.addfile(member, original_tar.extractfile(member))\n",
    "\t\t\t\t\tcurrent_split_size += member.size"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create datasets <a class=\"anchor\" id=\"head-1\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will be using the `COCO dataset`. `download_coco.sh` script from dataset prepare will be used to download and unzip the coco2017 dataset from [here](https://cocodataset.org/#download)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**If using custom dataset; it should follow this dataset structure**\n",
    "```\n",
    "DATA_DIR\n",
    "├── annotations.json\n",
    "├── images\n",
    "    ├── image_name_1.jpg\n",
    "    ├── image_name_2.jpg\n",
    "    ├── ...\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "dataset_to_be_used = \"default\" #FIXME7 #default/custom; default for the dataset used in this tutorial notebook; custom for a different dataset\n",
    "DATA_DIR = model_name # FIXME8\n",
    "os.environ['DATA_DIR']= DATA_DIR\n",
    "!mkdir -p $DATA_DIR"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Download dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if dataset_to_be_used == \"default\":\n",
    "    !bash ../dataset_prepare/coco/download_coco.sh $DATA_DIR\n",
    "    # Remove existing data\n",
    "    !rm -rf $DATA_DIR/train2017/images\n",
    "    !rm -rf $DATA_DIR/val2017/images\n",
    "    # Rearrange data in the required format\n",
    "    !mkdir -p $DATA_DIR/train2017/\n",
    "    !mkdir -p $DATA_DIR/val2017/\n",
    "    !mv $DATA_DIR/raw-data/train2017 $DATA_DIR/train2017/images\n",
    "    !mv $DATA_DIR/raw-data/annotations/instances_train2017.json $DATA_DIR/train2017/annotations.json\n",
    "    !mv $DATA_DIR/raw-data/val2017 $DATA_DIR/val2017/images\n",
    "    !mv $DATA_DIR/raw-data/annotations/instances_val2017.json $DATA_DIR/val2017/annotations.json\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Verify the downloaded dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!if [ ! -d $DATA_DIR/train2017/images ]; then echo 'Images folder not found'; else echo 'Found images folder';fi\n",
    "!if [ ! -f $DATA_DIR/train2017/annotations.json ]; then echo 'annotations file not found'; else echo 'Found annotations file';fi\n",
    "!if [ ! -d $DATA_DIR/val2017/images ]; then echo 'Images folder not found'; else echo 'Found images folder';fi\n",
    "!if [ ! -f $DATA_DIR/val2017/annotations.json ]; then echo 'annotations file not found'; else echo 'Found annotations file';fi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tar -C $DATA_DIR/train2017 -czf $DATA_DIR/coco_train.tar.gz images annotations.json\n",
    "!tar -C $DATA_DIR/val2017 -czf $DATA_DIR/coco_val.tar.gz images annotations.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset_path = f\"{DATA_DIR}/coco_train.tar.gz\"\n",
    "eval_dataset_path= f\"{DATA_DIR}/coco_val.tar.gz\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Create train dataset\n",
    "ds_type = \"instance_segmentation\"\n",
    "ds_format = \"coco\"\n",
    "data = json.dumps({\"type\":ds_type,\"format\":ds_format})\n",
    "\n",
    "endpoint = f\"{base_url}/datasets\"\n",
    "\n",
    "response = requests.post(endpoint,data=data,headers=headers)\n",
    "\n",
    "print(response)\n",
    "print(response.json())\n",
    "\n",
    "train_dataset_id = response.json()[\"id\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Update\n",
    "dataset_information = {\"name\":\"Train dataset\",\n",
    "                       \"description\":\"My train dataset\"}\n",
    "data = json.dumps(dataset_information)\n",
    "\n",
    "endpoint = f\"{base_url}/datasets/{train_dataset_id}\"\n",
    "\n",
    "response = requests.patch(endpoint, data=data, headers=headers)\n",
    "\n",
    "print(response)\n",
    "print(response.json())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Upload\n",
    "output_dir = os.path.join(os.path.dirname(os.path.abspath(train_dataset_path)), model_name, \"train\")\n",
    "split_tar_file(train_dataset_path, output_dir)\n",
    "for idx, tar_dataset_path in enumerate(os.listdir(output_dir)):\n",
    "    print(f\"Uploading {idx+1}/{len(os.listdir(output_dir))} tar split\")\n",
    "    files = [(\"file\",open(os.path.join(output_dir, tar_dataset_path),\"rb\"))]\n",
    "\n",
    "    endpoint = f\"{base_url}/datasets/{train_dataset_id}:upload\"\n",
    "\n",
    "    response = requests.post(endpoint, files=files, headers=headers)\n",
    "\n",
    "    print(response)\n",
    "    print(response.json())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create eval dataset\n",
    "data = json.dumps({\"type\":ds_type,\"format\":ds_format})\n",
    "\n",
    "endpoint = f\"{base_url}/datasets\"\n",
    "\n",
    "response = requests.post(endpoint,data=data,headers=headers)\n",
    "\n",
    "print(response)\n",
    "print(response.json())\n",
    "\n",
    "eval_dataset_id = response.json()[\"id\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Update\n",
    "dataset_information = {\"name\":\"Evaluation dataset\",\n",
    "                       \"description\":\"My eval dataset\"}\n",
    "data = json.dumps(dataset_information)\n",
    "\n",
    "endpoint = f\"{base_url}/datasets/{eval_dataset_id}\"\n",
    "\n",
    "response = requests.patch(endpoint, data=data, headers=headers)\n",
    "\n",
    "print(response)\n",
    "print(response.json())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_dir = os.path.join(os.path.dirname(os.path.abspath(eval_dataset_path)), model_name, \"eval\")\n",
    "split_tar_file(eval_dataset_path, output_dir)\n",
    "for idx, tar_dataset_path in enumerate(os.listdir(output_dir)):\n",
    "    print(f\"Uploading {idx+1}/{len(os.listdir(output_dir))} tar split\")\n",
    "    files = [(\"file\",open(os.path.join(output_dir, tar_dataset_path),\"rb\"))]\n",
    "\n",
    "    endpoint = f\"{base_url}/datasets/{eval_dataset_id}:upload\"\n",
    "\n",
    "    response = requests.post(endpoint, files=files, headers=headers)\n",
    "\n",
    "    print(response)\n",
    "    print(response.json())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### List the created datasets <a class=\"anchor\" id=\"head-2\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint = f\"{base_url}/datasets\"\n",
    "\n",
    "response = requests.get(endpoint, headers=headers)\n",
    "\n",
    "print(response)\n",
    "# print(response.json()) ## Uncomment for verbose list output\n",
    "print(\"id\\t\\t\\t\\t\\t type\\t\\t\\t format\\t\\t name\")\n",
    "for rsp in response.json():\n",
    "    print(rsp[\"id\"],\"\\t\",rsp[\"type\"],\"\\t\",rsp[\"format\"],\"\\t\\t\",rsp[\"name\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create an experiment <a class=\"anchor\" id=\"head-4\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "encode_key = \"tlt_encode\"\n",
    "checkpoint_choose_method = \"best_model\"\n",
    "data = json.dumps({\"network_arch\":model_name,\"encryption_key\":encode_key,\"checkpoint_choose_method\":checkpoint_choose_method})\n",
    "\n",
    "endpoint = f\"{base_url}/experiments\"\n",
    "\n",
    "response = requests.post(endpoint,data=data,headers=headers)\n",
    "\n",
    "print(response)\n",
    "print(response.json())\n",
    "experiment_id = response.json()[\"id\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### List experiments <a class=\"anchor\" id=\"head-5\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "endpoint = f\"{base_url}/experiments\"\n",
    "params = {\"network_arch\": model_name}\n",
    "response = requests.get(endpoint, params=params, headers=headers)\n",
    "\n",
    "print(response)\n",
    "# print(response.json()) ## Uncomment for verbose list output\n",
    "print(\"model id\\t\\t\\t     network architecture\")\n",
    "for rsp in response.json():\n",
    "    print(rsp[\"name\"],rsp[\"id\"],rsp[\"network_arch\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Assign train, eval datasets <a class=\"anchor\" id=\"head-6\"></a>\n",
    "\n",
    "- Note: make sure the order for train_datasets is [source ID, target ID]\n",
    "- eval_dataset is kept same as target for demo purposes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_information = {\"train_datasets\":[train_dataset_id],\n",
    "                       \"eval_dataset\":eval_dataset_id,\n",
    "                       \"inference_dataset\":eval_dataset_id}\n",
    "data = json.dumps(dataset_information)\n",
    "\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}\"\n",
    "\n",
    "response = requests.patch(endpoint, data=data, headers=headers)\n",
    "\n",
    "print(response)\n",
    "print(response.json())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Assign PTM <a class=\"anchor\" id=\"head-7\"></a>\n",
    "\n",
    "Search for the PTM on NGC for the Segmentation model chosen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# List all pretrained models for the chosen network architecture\n",
    "endpoint = f\"{base_url}/experiments\"\n",
    "params = {\"network_arch\": model_name}\n",
    "response = requests.get(endpoint, params=params, headers=headers)\n",
    "\n",
    "response_json = response.json()\n",
    "\n",
    "for rsp in response_json:\n",
    "    if \"encryption_key\" not in rsp.keys():\n",
    "        print(f'PTM Name: {rsp[\"name\"]}; PTM version: {rsp[\"version\"]}; NGC PATH: {rsp[\"ngc_path\"]}; Additional info: {rsp[\"additional_id_info\"]}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assigning pretrained models to different networks\n",
    "# From the output of previous cell make the appropriate changes to this map if you want to change the default PTM backbone.\n",
    "# Changing the default backbone here requires changing default spec/config during train/eval etc like for example\n",
    "# If you are changing the ptm to resnet34, then you have to modify the config key num_layers if it exists to 34 manually\n",
    "pretrained_map = {\"mal\" : \"mask_auto_label:trainable_v1.0\"}\n",
    "no_ptm_models = set([])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get pretrained model\n",
    "if model_name not in no_ptm_models:\n",
    "    endpoint = f\"{base_url}/experiments\"\n",
    "    params = {\"network_arch\": model_name}\n",
    "    response = requests.get(endpoint, params=params, headers=headers)\n",
    "\n",
    "    response_json = response.json()\n",
    "\n",
    "    # Search for ptm with given ngc path\n",
    "    ptm = []\n",
    "    for rsp in response_json:\n",
    "        if rsp[\"ngc_path\"].endswith(pretrained_map[model_name]):\n",
    "            ptm_id = rsp[\"id\"]\n",
    "            ptm = [ptm_id]\n",
    "            print(\"Metadata for model with requested NGC Path\")\n",
    "            print(rsp)\n",
    "            break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if model_name not in no_ptm_models:\n",
    "    ptm_information = {\"base_experiment\":ptm}\n",
    "    data = json.dumps(ptm_information)\n",
    "\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}\"\n",
    "\n",
    "    response = requests.patch(endpoint, data=data, headers=headers)\n",
    "\n",
    "    print(response)\n",
    "    print(response.json())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### View hyperparameters that are enabled for AutoML by default <a class=\"anchor\" id=\"head-8\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if automl_enabled:\n",
    "    # Get default spec schema\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/specs/train/schema\"\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    automl_specs = response.json()[\"automl_default_parameters\"]\n",
    "    print(json.dumps(automl_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set AutoML related configurations <a class=\"anchor\" id=\"head-9\"></a>\n",
    "Refer to these hyper-links to see the parameters supported by each network and add more parameters if necessary in addition to the default automl enabled parameters: [Mask RCNN](https://docs.nvidia.com/tao/tao-toolkit/text/tao_toolkit_api/api_action_specs.html#id39), [Unet](https://docs.nvidia.com/tao/tao-toolkit/text/tao_toolkit_api/api_action_specs.html#id69), [Segformer](https://docs.nvidia.com/tao/tao-toolkit/text/tao_toolkit_api/api_action_specs.html#id69)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if automl_enabled:\n",
    "    # Choose any metric that is present in the kpi dictionary present in the model's status.json. \n",
    "    # Example status.json for each model can be found in the respective section in NVIDIA TAO DOCS here: https://docs.nvidia.com/tao/tao-toolkit/text/model_zoo/cv_models/index.html\n",
    "    metric=\"kpi\"\n",
    "\n",
    "    additional_automl_parameters = [] #Refer to parameter list mentioned in the above links and add any extra parameter in addition to the default enabled ones\n",
    "    remove_default_automl_parameters = [] #Remove any hyperparameters that are enabled by default for AutoML\n",
    "\n",
    "    automl_information = {\"automl_enabled\":automl_enabled,\n",
    "                          \"automl_algorithm\":automl_algorithm,\n",
    "                          \"metric\":metric,\n",
    "                          \"automl_add_hyperparameters\":str(additional_automl_parameters),\n",
    "                          \"automl_remove_hyperparameters\":str(remove_default_automl_parameters)\n",
    "                         }\n",
    "    data = json.dumps(automl_information)\n",
    "\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}\"\n",
    "\n",
    "    response = requests.patch(endpoint, data=data, headers=headers)\n",
    "\n",
    "    print(response)\n",
    "    print(json.dumps(response.json(), sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Actions <a class=\"anchor\" id=\"head-10\"></a>\n",
    "\n",
    "For all actions:\n",
    "1. Get default spec schema and derive the default values\n",
    "2. Modify defaults if needed\n",
    "3. Post spec dictionary to the service\n",
    "4. Run model action\n",
    "5. Monitor job using retrieve\n",
    "6. Download results using job download endpoint (if needed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "job_map = {}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train <a class=\"anchor\" id=\"head-11\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Get default spec schema\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/specs/train/schema\"\n",
    "\n",
    "response = requests.get(endpoint, headers=headers)\n",
    "\n",
    "print(response)\n",
    "#print(response.json()) ## Uncomment for verbose schema\n",
    "train_specs = response.json()[\"default\"]\n",
    "print(json.dumps(train_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Override any of the parameters listed in the previous cell as required\n",
    "train_specs[\"gpu_ids\"] = [0]\n",
    "train_specs[\"train\"][\"num_epochs\"] = 5\n",
    "print(json.dumps(train_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run action\n",
    "parent = None\n",
    "action = \"train\"\n",
    "data = json.dumps({\"parent_job_id\":parent,\"action\":action, \"specs\":train_specs})\n",
    "\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/jobs\"\n",
    "\n",
    "response = requests.post(endpoint, data=data, headers=headers)\n",
    "\n",
    "print(response)\n",
    "print(response.json())\n",
    "\n",
    "job_map[\"train_\"+model_name] = response.json()\n",
    "print(job_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Monitor job status by repeatedly running this cell\n",
    "# For automl: Training times for different models benchmarked on 1 GPU V100 machine can be found here: https://docs.nvidia.com/tao/tao-toolkit/text/automl/automl.html#results-of-automl-experiments\n",
    "\n",
    "job_id = job_map['train_'+model_name]\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}\"\n",
    "\n",
    "while True:    \n",
    "    clear_output(wait=True)\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    print(response)\n",
    "    print(json.dumps(response.json(), sort_keys=True, indent=4))\n",
    "    if response.json().get(\"status\") in [\"Done\",\"Error\", \"Canceled\"] or response.status_code not in (200,201):\n",
    "        break\n",
    "    time.sleep(15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## To Stop an AutoML JOB\n",
    "#    1. Stop the 'Monitor job status by repeatedly running this cell' cell (the cell right before this cell) manually\n",
    "#    2. Uncomment the snippet in the next cell and run the cell"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# if automl_enabled:\n",
    "#     job_id = job_map['train_'+model_name]\n",
    "#     endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}:cancel\"\n",
    "\n",
    "#     response = requests.post(endpoint, headers=headers)\n",
    "\n",
    "#     print(response)\n",
    "#     print(response.json())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Resume AutoML"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Uncomment the below snippet if you want to resume an already stopped AutoML job and then run the 'Monitor job status by repeatedly running this cell' cell above (4th cell above from this cell)\n",
    "# if automl_enabled:\n",
    "#     job_id = job_map['train_'+model_name]\n",
    "#     endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}:resume\"\n",
    "\n",
    "#     data = json.dumps({\"parent_job_id\":parent,\"specs\":train_specs})\n",
    "#     response = requests.post(endpoint, data=data, headers=headers)\n",
    "\n",
    "#     print(response)\n",
    "#     print(response.json())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Example to list the files of the executed train job\n",
    "job_id = job_map[\"train_\"+model_name]\n",
    "endpoint = f'{base_url}/experiments/{experiment_id}/jobs/{job_id}:list_files'\n",
    "\n",
    "response = requests.get(endpoint, headers=headers)\n",
    "print(json.dumps(response.json(), sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "## Patch the model with proper metric before training to run this cell; By default loss is used, but some models dont log the parameter under the name 'loss'\n",
    "\n",
    "# # Download selective job contents once the above job shows \"Done\" status\n",
    "# # Example to download selective files of train job (Note: will take time)\n",
    "# job_id = job_map[\"train_\"+model_name]\n",
    "# endpoint = f'{base_url}/experiments/{experiment_id}/jobs/{job_id}:download_selective_files'\n",
    "\n",
    "# file_lists = [] # Choose file names from the previous cell where all the files for this job were listed\n",
    "# best_model = False # Enable this to download the checkpoint of the best performing model w.r.t to the metric chosen before starting training\n",
    "# latest_model = True # Enable this to download the latest checkpoint of the training job; Disable best_model to use latest_model\n",
    "\n",
    "# params = {\"file_lists\": file_lists, \"best_model\": best_model, \"latest_model\": latest_model}\n",
    "\n",
    "# # Save\n",
    "# temptar = f'{job_id}.tar.gz'\n",
    "# with requests.get(endpoint, headers=headers, params=params, stream=True) as r:\n",
    "#     r.raise_for_status()\n",
    "#     with open(temptar, 'wb') as f:\n",
    "#         for chunk in r.iter_content(chunk_size=8192):\n",
    "#             f.write(chunk)\n",
    "\n",
    "# print(\"Untarring\")\n",
    "# # Untar to destination\n",
    "# tar_command = f'tar -xvf {temptar} -C {workdir}/'\n",
    "# os.system(tar_command)\n",
    "# os.remove(temptar)\n",
    "# print(f\"Results at {workdir}/{job_id}\")\n",
    "# model_downloaded_path = f\"{workdir}/{job_id}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Downloading train job takes a longer time, uncomment this cell if you want to still proceed\n",
    "if download_jobs:\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}\"\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    expected_file_size = response.json().get(\"job_tar_stats\", {}).get(\"file_size\")\n",
    "    print(\"expected_file_size: \", expected_file_size)\n",
    "\n",
    "    !python3 -m pip install tqdm\n",
    "    from tqdm import tqdm\n",
    "\n",
    "    endpoint = f'{base_url}/experiments/{experiment_id}/jobs/{job_id}:download'\n",
    "    temptar = f'{job_id}.tar.gz'\n",
    "\n",
    "    with tqdm(total=expected_file_size, unit='B', unit_scale=True) as progress_bar:\n",
    "        while True:\n",
    "            # Check if the file already exists\n",
    "            headers_download_job = dict(headers)\n",
    "            if os.path.exists(temptar):\n",
    "                # Get the current file size\n",
    "                file_size = os.path.getsize(temptar)\n",
    "                print(f\"File size of dowloaded content until now is {file_size}\")\n",
    "\n",
    "                # If the file size matches the expected size, break out of the loop\n",
    "                if file_size >= (expected_file_size-1):\n",
    "                    print(\"Download completed successfully.\")\n",
    "                    print(\"Untarring\")\n",
    "                    # Untar to destination\n",
    "                    tar_command = f'tar -xf {temptar} -C {workdir}/'\n",
    "                    os.system(tar_command)\n",
    "                    os.remove(temptar)\n",
    "                    print(f\"Results at {workdir}/{job_id}\")\n",
    "                    model_downloaded_path = f\"{workdir}/{job_id}\"\n",
    "                    break\n",
    "\n",
    "                # Set the headers to resume the download from where it left off\n",
    "                headers_download_job['Range'] = f'bytes={file_size}-'\n",
    "            # Open the file for writing in binary mode\n",
    "            with open(temptar, 'ab') as f:\n",
    "                try:\n",
    "                    response = requests.get(endpoint, headers=headers_download_job, stream=True)\n",
    "                    print(response)\n",
    "                    # Check if the request was successful\n",
    "                    if response.status_code in [200, 206]:\n",
    "                        # Iterate over the content in chunks\n",
    "                        for chunk in response.iter_content(chunk_size=1024):\n",
    "                            if chunk:\n",
    "                                # Write the chunk to the file\n",
    "                                f.write(chunk)\n",
    "                                # Flush and sync the file to disk\n",
    "                                f.flush()\n",
    "                                os.fsync(f.fileno())\n",
    "                            progress_bar.update(len(chunk))\n",
    "                    else:\n",
    "                        print(f\"Failed to download file. Status code: {response.status_code}\")\n",
    "                except requests.exceptions.RequestException as e:\n",
    "                    print(\"Connection interrupted during download, resuming download from breaking point\")\n",
    "                    time.sleep(5)  # Sleep for a while before retrying the request\n",
    "                    continue  # Continue the loop to retry the request"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View the checkpoints generated for the training job and for automl jobs, in addition view: best performing model's config and the results of all automl experiments\n",
    "\n",
    "if download_jobs:\n",
    "    if automl_enabled:\n",
    "        !python3 -m pip install pandas==1.5.1\n",
    "        import pandas as pd\n",
    "        model_downloaded_path = f\"{model_downloaded_path}/best_model\"\n",
    "        assert glob.glob(f\"{model_downloaded_path}/*.protobuf\") or glob.glob(f\"{model_downloaded_path}/*.yaml\")\n",
    "\n",
    "    assert os.path.exists(model_downloaded_path)\n",
    "    assert (glob.glob(model_downloaded_path + \"/**/*.tlt\", recursive=True) + glob.glob(model_downloaded_path + \"/**/*.hdf5\", recursive=True) + glob.glob(model_downloaded_path + \"/**/*.pth\", recursive=True))\n",
    "\n",
    "    if os.path.exists(model_downloaded_path):        \n",
    "        #List the binary model file\n",
    "        print(\"\\nCheckpoints for the training experiment\")\n",
    "        if os.path.exists(model_downloaded_path+\"/train/weights\") and len(os.listdir(model_downloaded_path+\"/train/weights\")) > 0:\n",
    "            print(f\"Folder: {model_downloaded_path}/train/weights\")\n",
    "            print(\"Files:\", os.listdir(model_downloaded_path+\"/train/weights\"))\n",
    "        elif os.path.exists(model_downloaded_path+\"/weights\") and len(os.listdir(model_downloaded_path+\"/weights\")) > 0:\n",
    "            print(f\"Folder: {model_downloaded_path}/weights\")\n",
    "            print(\"Files:\", os.listdir(model_downloaded_path+\"/weights\"))\n",
    "        else:\n",
    "            print(f\"Folder: {model_downloaded_path}\")\n",
    "            print(\"Files:\", os.listdir(model_downloaded_path))\n",
    "\n",
    "        if automl_enabled:\n",
    "            assert glob.glob(f\"{model_downloaded_path}/*.protobuf\") or glob.glob(f\"{model_downloaded_path}/*.yaml\")\n",
    "            experiment_artifacts = json.load(open(f\"{model_downloaded_path}/controller.json\",\"r\"))\n",
    "            data_frame = pd.DataFrame(experiment_artifacts)\n",
    "            # Print experiment id/number and the corresponding result\n",
    "            print(\"\\nResults of all experiments\")\n",
    "            with pd.option_context('display.max_rows', None, 'display.max_columns', None, 'display.max_colwidth', None):\n",
    "                print(data_frame[[\"id\",\"result\"]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluate <a class=\"anchor\" id=\"head-12\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Get model handler parameters\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}\"\n",
    "response = requests.get(endpoint, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "assert response.json()\n",
    "\n",
    "model_parameters = response.json()\n",
    "update_checkpoint_choosing = {}\n",
    "update_checkpoint_choosing[\"checkpoint_choose_method\"] = model_parameters[\"checkpoint_choose_method\"]\n",
    "update_checkpoint_choosing[\"checkpoint_epoch_number\"] = model_parameters[\"checkpoint_epoch_number\"]\n",
    "print(update_checkpoint_choosing)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Change the method by which checkpoint from the parent action is chosen, when parent action is a train/retrain action.\n",
    "# Example for evaluate action below, can be applied in the same way for other actions too\n",
    "update_checkpoint_choosing[\"checkpoint_choose_method\"] = \"latest_model\" # Choose between best_model/latest_model/from_epoch_number\n",
    "# If from_epoch_number is chosen then assign the epoch number to the dictionary key in the format 'from_epoch_number{train_job_id}'\n",
    "# update_checkpoint_choosing[\"checkpoint_epoch_number\"][\"from_epoch_number_28a2754e-50ef-43a8-9733-98913776dd90\"] = 3\n",
    "data = json.dumps(update_checkpoint_choosing)\n",
    "\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}\"\n",
    "\n",
    "response = requests.patch(endpoint, data=data, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "print(response)\n",
    "print(json.dumps(response.json(), sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Get default spec schema\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/specs/evaluate/schema\"\n",
    "\n",
    "response = requests.get(endpoint, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "print(response)\n",
    "#print(response.json()) ## Uncomment for verbose schema\n",
    "assert \"default\" in response.json().keys()\n",
    "eval_specs = response.json()[\"default\"]\n",
    "print(json.dumps(eval_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Apply changes to the specs if required\n",
    "print(json.dumps(eval_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run action\n",
    "parent = job_map[\"train_\" + model_name]\n",
    "action = \"evaluate\"\n",
    "data = json.dumps({\"parent_job_id\":parent,\"action\":action,\"specs\":eval_specs})\n",
    "\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/jobs\"\n",
    "\n",
    "response = requests.post(endpoint, data=data, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "assert response.json()\n",
    "\n",
    "print(response)\n",
    "print(response.json())\n",
    "\n",
    "job_map[\"evaluate_\" + model_name] = response.json()\n",
    "print(job_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Monitor job status by repeatedly running this cell\n",
    "job_id = job_map[\"evaluate_\" + model_name]\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}\"\n",
    "\n",
    "while True:    \n",
    "    clear_output(wait=True)\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    print(response)\n",
    "    print(response.json())\n",
    "    assert \"status\" in response.json().keys() and response.json().get(\"status\") != \"Error\"\n",
    "    if response.json().get(\"status\") in [\"Done\",\"Error\", \"Canceled\"] or response.status_code not in (200,201):\n",
    "        break\n",
    "    time.sleep(15)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TAO inference <a class=\"anchor\" id=\"head-20\"></a>\n",
    "\n",
    "- Run inference on a set of images using the .tlt model created at train step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Get default spec schema\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/specs/inference/schema\"\n",
    "\n",
    "response = requests.get(endpoint, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "print(response)\n",
    "#print(response.json()) ## Uncomment for verbose schema\n",
    "assert \"default\" in response.json().keys()\n",
    "tao_inference_specs = response.json()[\"default\"]\n",
    "print(json.dumps(tao_inference_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Apply changes to specs if necessary\n",
    "print(json.dumps(tao_inference_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run action\n",
    "parent = job_map[\"train_\" + model_name]\n",
    "action = \"inference\"\n",
    "data = json.dumps({\"parent_job_id\":parent,\"action\":action,\"specs\":tao_inference_specs})\n",
    "\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/jobs\"\n",
    "\n",
    "response = requests.post(endpoint, data=data, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "assert response.json()\n",
    "\n",
    "print(response)\n",
    "print(response.json())\n",
    "\n",
    "job_map[\"inference_tlt_\" + model_name] = response.json()\n",
    "print(job_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Monitor job status by repeatedly running this cell\n",
    "job_id = job_map[\"inference_tlt_\" + model_name]\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}\"\n",
    "\n",
    "while True:\n",
    "    clear_output(wait=True)\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    print(response)\n",
    "    print(response.json())\n",
    "    assert \"status\" in response.json().keys() and response.json().get(\"status\") != \"Error\"\n",
    "    if response.json().get(\"status\") in [\"Done\",\"Error\", \"Canceled\"] or response.status_code not in (200,201):\n",
    "        break\n",
    "    time.sleep(15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download job contents once the above job shows \"Done\" status\n",
    "if download_jobs:\n",
    "    job_id = job_map[\"inference_tlt_\" + model_name]\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}\"\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    expected_file_size = response.json().get(\"job_tar_stats\", {}).get(\"file_size\")\n",
    "    print(\"expected_file_size: \", expected_file_size)\n",
    "\n",
    "    !python3 -m pip install tqdm\n",
    "    from tqdm import tqdm\n",
    "\n",
    "    endpoint = f'{base_url}/experiments/{experiment_id}/jobs/{job_id}:download'\n",
    "    temptar = f'{job_id}.tar.gz'\n",
    "\n",
    "    with tqdm(total=expected_file_size, unit='B', unit_scale=True) as progress_bar:\n",
    "        while True:\n",
    "            # Check if the file already exists\n",
    "            headers_download_job = dict(headers)\n",
    "            if os.path.exists(temptar):\n",
    "                # Get the current file size\n",
    "                file_size = os.path.getsize(temptar)\n",
    "                print(f\"File size of dowloaded content until now is {file_size}\")\n",
    "\n",
    "                # If the file size matches the expected size, break out of the loop\n",
    "                if file_size >= (expected_file_size-1):\n",
    "                    print(\"Download completed successfully.\")\n",
    "                    print(\"Untarring\")\n",
    "                    # Untar to destination\n",
    "                    tar_command = f'tar -xf {temptar} -C {workdir}/'\n",
    "                    os.system(tar_command)\n",
    "                    os.remove(temptar)\n",
    "                    print(f\"Results at {workdir}/{job_id}\")\n",
    "                    inference_out_path = f\"{workdir}/{job_id}\"\n",
    "                    break\n",
    "\n",
    "                # Set the headers to resume the download from where it left off\n",
    "                headers_download_job['Range'] = f'bytes={file_size}-'\n",
    "            # Open the file for writing in binary mode\n",
    "            with open(temptar, 'ab') as f:\n",
    "                try:\n",
    "                    response = requests.get(endpoint, headers=headers_download_job, stream=True)\n",
    "                    print(response)\n",
    "                    # Check if the request was successful\n",
    "                    if response.status_code in [200, 206]:\n",
    "                        # Iterate over the content in chunks\n",
    "                        for chunk in response.iter_content(chunk_size=1024):\n",
    "                            if chunk:\n",
    "                                # Write the chunk to the file\n",
    "                                f.write(chunk)\n",
    "                                # Flush and sync the file to disk\n",
    "                                f.flush()\n",
    "                                os.fsync(f.fileno())\n",
    "                            progress_bar.update(len(chunk))\n",
    "                    else:\n",
    "                        print(f\"Failed to download file. Status code: {response.status_code}\")\n",
    "                except requests.exceptions.RequestException as e:\n",
    "                    print(\"Connection interrupted during download, resuming download from breaking point\")\n",
    "                    time.sleep(5)  # Sleep for a while before retrying the request\n",
    "                    continue  # Continue the loop to retry the request"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Inference output must be here\n",
    "if download_jobs:\n",
    "    !ls {inference_out_path}/inference.json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Delete model <a class=\"anchor\" id=\"head-21\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint = f\"{base_url}/experiments/{experiment_id}\"\n",
    "\n",
    "response = requests.delete(endpoint,headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "print(response)\n",
    "print(response.json())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Delete dataset <a class=\"anchor\" id=\"head-21\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Delete train dataset <a class=\"anchor\" id=\"head-21\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint = f\"{base_url}/datasets/{train_dataset_id}\"\n",
    "\n",
    "response = requests.delete(endpoint,headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "print(response)\n",
    "print(response.json())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Delete val dataset <a class=\"anchor\" id=\"head-21\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint = f\"{base_url}/datasets/{eval_dataset_id}\"\n",
    "\n",
    "response = requests.delete(endpoint,headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "print(response)\n",
    "print(response.json())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
